import numpy as np
from itertools import *
from functools import *
import operator
from scipy.stats import norm


# 最大值函数
def I(x: np.ndarray):
    t = np.zeros([x.shape[0], 2])
    t[:, 0] = x
    return np.max(t, axis=1)


# 期望函数
def E(x: np.ndarray, p: np.ndarray):
    return np.dot(x, p)


# 依照正态分布计算概率值
def p_norm(x: np.ndarray, mu, sigma):
    init_arr = norm.cdf(x, mu, sigma)
    res_arr = np.zeros(x.shape[0])
    res_arr[0] = init_arr[0]
    res_arr[1:] = np.diff(init_arr)
    res_arr[-1] = 1 - init_arr[-1]
    return res_arr


# Variable
T = 3
alpha = 0.9
# Dm
mu_m = 2
sigma_m = 0.1
D_m = np.array(list(range(0, 6)))
p_D_m = p_norm(D_m, mu_m, sigma_m)
# Dt
mu_t = 5
sigma_t = 0.5
D_t = np.array(list(range(0, 10)))
p_D_t = p_norm(D_t, mu_t, sigma_t)
p = 10
c_e = 20
c_m = 2
h_e = 5
h_m = 1
theta = h_m / (1 - alpha)
b = 2

initial = True

# 多线程优化模块
from concurrent.futures import ProcessPoolExecutor, as_completed


def grid_search_ex(iter_grid, func, params: tuple) -> list:
    tasks = []
    pool = ProcessPoolExecutor(max_workers=6)

    for x, y in iter_grid:
        tasks.append(pool.submit(func, *params, x, y))

    return [x.result() for x in as_completed(tasks)]


def linear_search_ex(iter_arr, func, params: tuple) -> list:
    tasks = []
    pool = ProcessPoolExecutor(max_workers=6)

    for x in iter_arr:
       tasks.append(pool.submit(func, *params, x))

    return [x.result() for x in as_completed(tasks)]


# 以下为 t = T 时的代码, 可以看出, 与任意 t 时刻的不同之处在于 psi 处的 f 函数取值为 0(初值条件)
def psi_1T(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z):
    return h_e * E(I(v_1_bar - e - D_t), p_D_t) \
           + b * E(I(D_t - v_1_bar + e), p_D_t) \
           + c_e * (v_1_bar - e) \
           + c_m * (z - v_4_bar) \
           + h_m * (v_4_bar - v_1_bar) \
           + theta * e


def psi_2T(v_2_bar, v_3_bar, v_4_bar, y, z):
    return h_e * E(I(y - D_t), p_D_t) \
           + b * E(I(D_t - y), p_D_t) \
           + c_e * y \
           + c_m * (z - v_4_bar) \
           + h_m * (v_4_bar - y)


def psi_3T(v_3_bar, v_4_bar, y, z):
    return h_e * E(I(y - D_t), p_D_t) \
           + b * E(I(D_t - y), p_D_t) \
           + c_e * y \
           + c_m * (z - v_4_bar) \
           + h_m * (v_4_bar - y)


def psi_4T(v_4_bar, y, z):
    return h_e * E(I(y - D_t), p_D_t) \
           + b * E(I(D_t - y), p_D_t) \
           + c_e * y \
           + c_m * (z - v_4_bar) \
           + h_m * (v_4_bar - y)


def phi_1T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    e_range = range(0, v_1_bar - v_0_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    e_z_grid = product(e_range, z_range)
    psi_1T_values = grid_search_ex(e_z_grid, psi_1T, (v_1_bar, v_2_bar, v_3_bar, v_4_bar))
    return - c_e * v_0_bar + np.min(psi_1T_values)


def phi_2T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_1_bar, v_2_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_2T_values = grid_search_ex(y_z_grid, psi_2T, (v_2_bar, v_3_bar, v_4_bar))
    return - c_e * v_0_bar + np.min(psi_2T_values)


def phi_3T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_2_bar, v_3_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_3T_values = grid_search_ex(y_z_grid, psi_3T, (v_3_bar, v_4_bar))
    return - c_e * v_0_bar + np.min(psi_3T_values)


def phi_4T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_3_bar, v_4_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_4T_values = grid_search_ex(y_z_grid, psi_4T, (v_4_bar,))
    return - c_e * v_0_bar + np.min(psi_4T_values)


def g_T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    return np.min([
        phi_1T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_2T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_3T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_4T(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar)
    ])


def f_T(v_0, v_1, v_2, v_3, v_4):
    _D_m = []
    for D_m_i in D_m:
        exp_1_value = np.min([p * (D_m_i - a) + g_T(v_0, v_1 - a, v_2 - a, v_3 - a, v_4 - a) for a in
                              range(0, np.min([D_m_i, v_1 - v_0]) + 1)])
        exp_2_value = np.min([p * (D_m_i - a) + g_T(v_0, v_0, v_2 - a, v_3 - a, v_4 - a) for a in
                              range(0, np.min([D_m_i, v_2 - v_0]) + 1)])
        exp_3_value = np.min(
            [p * (D_m_i - a) + g_T(v_0, v_0, v_0, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_3 - v_0]) + 1)])
        exp_4_value = np.min(
            [p * (D_m_i - a) + g_T(v_0, v_0, v_0, v_0, v_4 - a) for a in range(0, np.min([D_m_i, v_4 - v_0]) + 1)])
        _D_m.append(np.min([exp_1_value, exp_2_value, exp_3_value, exp_4_value]))
    return E(_D_m, p_D_m)


# 以下为任意 t 时刻的代码
def psi_1t(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z):
    global initial
    val = psi_1T(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z)
    if initial:
        return val

    _D_t = []
    for D_t_i in D_t:
        _D_t.append(
            f_T(v_1_bar - e - D_t_i, v_2_bar - e - D_t_i, v_3_bar - e - D_t_i, v_4_bar - e - D_t_i, z - e - D_t_i))
    return val + alpha * E(_D_t, p_D_t)


def psi_2t(v_2_bar, v_3_bar, v_4_bar, y, z):
    global initial
    val = psi_2T(v_2_bar, v_3_bar, v_4_bar, y, z)
    if initial:
        return val

    _D_t = []
    for D_t_i in D_t:
        _D_t.append(f_T(y - D_t_i, v_2_bar - D_t_i, v_3_bar - D_t_i, v_4_bar - D_t_i, z - D_t_i))
    return val + alpha * E(_D_t, p_D_t)


def psi_3t(v_3_bar, v_4_bar, y, z):
    global initial
    val = psi_3T(v_3_bar, v_4_bar, y, z)
    if initial:
        return val

    _D_t = []
    for D_t_i in D_t:
        _D_t.append(f_T(y - D_t_i, y - D_t_i, v_3_bar - D_t_i, v_4_bar - D_t_i, z - D_t_i))
    return val + alpha * E(_D_t, p_D_t)


def psi_4t(v_4_bar, y, z):
    global initial
    val = psi_4T(v_4_bar, y, z)
    if initial:
        return val

    _D_t = []
    for D_t_i in D_t:
        _D_t.append(f_T(y - D_t_i, y - D_t_i, y - D_t_i, v_4_bar - D_t_i, z - D_t_i))
    return val + alpha * E(_D_t, p_D_t)


def phi_1t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    e_range = range(0, v_1_bar - v_0_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    e_z_grid = product(e_range, z_range)
    psi_1t_values = [psi_1t(v_1_bar, v_2_bar, v_3_bar, v_4_bar, e, z) for e, z in e_z_grid]
    return - c_e * v_0_bar + np.min(psi_1t_values)


def phi_2t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_1_bar, v_2_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_2t_values = [psi_2t(v_2_bar, v_3_bar, v_4_bar, y, z) for y, z in y_z_grid]
    return - c_e * v_0_bar + np.min(psi_2t_values)


def phi_3t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_2_bar, v_3_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_3t_values = [psi_3t(v_3_bar, v_4_bar, y, z) for y, z in y_z_grid]
    return - c_e * v_0_bar + np.min(psi_3t_values)


def phi_4t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    y_range = range(v_3_bar, v_4_bar + 1)
    z_range = range(v_4_bar, v_4_bar + 21)
    y_z_grid = product(y_range, z_range)
    psi_4t_values = [psi_4t(v_4_bar, y, z) for y, z in y_z_grid]
    return - c_e * v_0_bar + np.min(psi_4t_values)


def g_t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar):
    return np.min([
        phi_1t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_2t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_3t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar),
        phi_4t(v_0_bar, v_1_bar, v_2_bar, v_3_bar, v_4_bar)
    ])


def f_t(v_0, v_1, v_2, v_3, v_4):
    _D_m = []
    for D_m_i in D_m:
        exp_1_value = np.min([p * (D_m_i - a) + g_t(v_0, v_1 - a, v_2 - a, v_3 - a, v_4 - a) for a in
                              range(0, np.min([D_m_i, v_1 - v_0]) + 1)])
        exp_2_value = np.min([p * (D_m_i - a) + g_t(v_0, v_0, v_2 - a, v_3 - a, v_4 - a) for a in
                              range(0, np.min([D_m_i, v_2 - v_0]) + 1)])
        exp_3_value = np.min(
            [p * (D_m_i - a) + g_t(v_0, v_0, v_0, v_3 - a, v_4 - a) for a in range(0, np.min([D_m_i, v_3 - v_0]) + 1)])
        exp_4_value = np.min(
            [p * (D_m_i - a) + g_t(v_0, v_0, v_0, v_0, v_4 - a) for a in range(0, np.min([D_m_i, v_4 - v_0]) + 1)])
        _D_m.append(np.min([exp_1_value, exp_2_value, exp_3_value, exp_4_value]))
    return E(_D_m, p_D_m)


if __name__ == "__main__":
    initial = False
    print(f_t(1, 2, 3, 4, 5))
